#include <stdio.h>
#include <string.h>
#include <malloc.h>
#include <stdlib.h>
#include <math.h>
#include <cuda_runtime.h>
#include "cutrace.h"

//const double msmi = 2.31/100.0;

double Bfr[Nx*Ny];
double Bfz[Nx*Ny];
double Bfphi[Nx*Ny];
double Apsi[Nx*Ny];

__device__ vec cudacross(vec a, vec b)
{
    vec c = {0.0,0.0,0.0};
    c.x = a.y*b.z - a.z*b.y;
    c.y = a.z*b.x - a.x*b.z;
    c.z = a.x*b.y - a.y*b.x;
    return c;
}
__device__ vec cudaadd(vec a, vec b)
{
    vec c = {0.0,0.0,0.0};
    c.x = a.x+b.x;
    c.y = a.y+b.y;
    c.z = a.z+b.z;
    return c;
}
__device__ vec cudamul(double a, vec b)
{
    vec c = {0.0,0.0,0.0};
    c.x = a*b.x;
    c.y = a*b.y;
    c.z = a*b.z;
    return c;
}
__device__ double cudadot(vec a, vec b)
{
    return a.x*b.x + a.y*b.y + a.z*b.z;
}
// va represent vector array
__global__ void cupush(vec* vap, vec*vav, vec* vaB, double dt, double qmr)
{
    int i = blockDim.x*blockIdx.x + threadIdx.x;
    double step = dt*qmr/2.0;
    vec vs = {step*vaB[i].x, step*vaB[i].y, step*vaB[i].z};
    double fctr = 2.0/(1.0+cudadot(vs, vs));  // determinant factor
    vec vfor = cudacross(vav[i], vs);
    vec vback = cudacross(vfor, vs);
    vav[i] = cudaadd(vav[i], cudamul(fctr, cudaadd(vfor, vback)));
    vap[i] = cudaadd(vap[i], cudamul(dt, vav[i]));
}

/* Lagrange interpolation of second order*/
__global__ void cudainterp(vec* vap, double* valuarr, double* field)
{
    int i = blockDim.x*blockIdx.x + threadIdx.x;
    double pr = sqrt(vap[i].x*vap[i].x + vap[i].y*vap[i].y);
    double pz = vap[i].z;
    if (pr > Rmax || pr < Rmin || pz > Zmax || pz < Zmin)
        valuarr[i] = 0.0;
    int rg = (int)((pr - Rmin)/dr);
    int zg = (int)((pz - Zmin)/dz);
    double ofr = (pr - (double)rg * dr - Rmin)/dr;
    double ofz = (pz - (double)zg * dz - Zmin)/dz;
    double ret = 0.0;
    double coefrs[4] = {ofr*(ofr-1.0)*(ofr-2.0)/-6.0, (ofr+1.0)*(ofr-1.0)*(ofr-2.0)/2.0,
    (ofr+1.0)*ofr*(ofr-2.0)/-2.0, (ofr+1.0)*ofr*(ofr-1.0)/6.0};
    double coefzs[4] = {ofz*(ofz-1.0)*(ofz-2.0)/-6.0, (ofz+1.0)*(ofz-1.0)*(ofz-2.0)/2.0,
    (ofz+1.0)*ofz*(ofz-2.0)/-2.0, (ofz+1.0)*ofz*(ofz-1.0)/6.0};
    for(int i = 0; i < 4; i++){
    for(int j = 0; j < 4; j++)
    {
        ret += coefrs[i]*coefzs[j]*field[Ny*(rg+i-1)+zg+j-1];
    }}
    valuarr[i] = ret;
}

__global__ void cudavectorfieldform(vec* vaB, vec*vap, double* Bra, double* Bza, double* Bphia)
{
    int i = blockDim.x*blockIdx.x + threadIdx.x;
    vec vp = vap[i];
    double pphi = atan(vp.y/vp.x);
    pphi += pi*(vp.x<0.0);
    pphi += 2*pi*(vp.y < 0.0 && vp.x > 0.0);
    vec c = {0.0,0.0, Bza[i]};
    c.x = Bra[i]*cos(pphi) - Bphia[i]*sin(pphi);
    c.y = Bra[i]*sin(pphi) + Bphia[i]*cos(pphi); 
    vaB[i] = c;    
}

// check cross midplane using cuda, vagdcp: vector array of guiding center position
__global__ void cudamdp(vec* vap, vec* vav, vec* vaB, vec* vagdcp, int* vstart, double* vfreq, int time)
{
    int i = blockDim.x*blockIdx.x + threadIdx.x;
    vec vp = vap[i];
    vec vB = vaB[i];
    vec vv = vav[i];
    vec gdc = cudaadd(vp, cudamul(1.0/qmr/cudadot(vB,vB),cudacross(vv, vB)));
    if ((gdc.z > 0)^(vagdcp[i].z > 0) && time > 1000)
    {
        if(vstart[i] == 0)
            vstart[i] = time;
        else if(vfreq[i] == 0.0)
            vfreq[i] = 100000.0/(2.0*(time-vstart[i]));
    }
    vagdcp[i]= gdc;
    return;
}

int main(int argc, char* argv[]){
    
    int nptc = 32768;
    int Ncycle = 32;
    bool usegpu = true;
    if (argc == 2) {
    if (*argv[1] == 'c')
    { usegpu = false; printf("using cpu\n");} 
    else {printf("using gpu\n");} } else {printf("using gpu\n");}
    
    loadfield(Apsi,"Apsi.csv");
    loadfield(Bfr,"Br.csv");
    loadfield(Bfz,"Bz.csv");
    loadfield(Bfphi,"Bphi.csv");
    
    double dt = 1.0/Ncycle/qmr;

    vec* vparr = (vec*)malloc(sizeof(vec)*nptc);
    vec* vvarr = (vec*)malloc(sizeof(vec)*nptc);
    vec* vBarr = (vec*)malloc(sizeof(vec)*nptc);
    vec* vgdcarr = (vec*)malloc(sizeof(vec)*nptc);
    int* vstart = (int*)malloc(sizeof(int)*nptc);
    double* vfreq = (double*)malloc(sizeof(double)*nptc);

    initdata_passing(vparr, vvarr, nptc);
    for (int i = 0; i < nptc; i++) {vstart[i] = 0; vfreq[i] = 0.0;}

    for (int i = 0; i < nptc; i += 4*655)
    {
        diagnose(vparr[i], vvarr[i],'c');   
    }
    printf("\n");
    //printf("%lf, %lf, %p\n", Bfphi[20*Ny+20], *(Bfphi+20*Ny+20), Bfphi);

    if (usegpu == true)
    {
        int nthread = 1024;
        int nblock = 32;
        vec* cuvparr, *cuvvarr, *cuvBarr, *cuvgdc;
        double *cuBrf, *cuBzf, *cuBphif, *cuBra, *cuBza, *cuBphia;
        double *cufreq;
        int* custep;

        cudaMalloc((double **)&cuBra, sizeof(double)*nptc);
        cudaMalloc((double **)&cuBza, sizeof(double)*nptc);
        cudaMalloc((double **)&cuBphia, sizeof(double)*nptc);

        cudaMalloc((vec**)&cuvparr, sizeof(vec)*nptc);
        cudaMalloc((vec**)&cuvvarr, sizeof(vec)*nptc);
        cudaMalloc((vec**)&cuvBarr, sizeof(vec)*nptc);
        cudaMalloc((void**)&cuvgdc, sizeof(vec)*nptc);
        cudaMalloc((void**)&custep, sizeof(int)*nptc);
        cudaMalloc((void**)&cufreq, sizeof(double)*nptc);

        cudaMalloc((double **)&cuBrf, sizeof(double)*Nx*Ny);
        cudaMalloc((double **)&cuBzf, sizeof(double)*Nx*Ny);
        cudaMalloc((double **)&cuBphif, sizeof(double)*Nx*Ny);

        cudaMemcpy(cuBrf, Bfr, sizeof(double)*Nx*Ny, cudaMemcpyHostToDevice);
        cudaMemcpy(cuBzf, Bfz, sizeof(double)*Nx*Ny, cudaMemcpyHostToDevice);
        cudaMemcpy(cuBphif, Bfphi, sizeof(double)*Nx*Ny, cudaMemcpyHostToDevice);
        cudaMemcpy(cuvparr, vparr, sizeof(vec)*nptc, cudaMemcpyHostToDevice);
        cudaMemcpy(cuvvarr, vvarr, sizeof(vec)*nptc, cudaMemcpyHostToDevice);
        cudaMemcpy(custep, vstart, sizeof(int)*nptc, cudaMemcpyHostToDevice);
        cudaMemcpy(cufreq, vfreq, sizeof(double)*nptc, cudaMemcpyHostToDevice);

        bool interptest = false;
        if(interptest == true)
        {   
            double *cuvaluearr;
            double* valuearr = (double*)malloc(sizeof(double)*1024);
            cudaMalloc((double**)&cuvaluearr, sizeof(double)*1024);
            cudainterp<<<1,1024>>>(cuvparr, cuvaluearr, cuBphif);
            cudaMemcpy(valuearr,cuvaluearr, sizeof(double)*1024, cudaMemcpyDeviceToHost);
            for(int i = 0; i < 1024; i += 363)
            {
                printf("interp cpu: %.10lf ",
                interp(sqrt(pow(vparr[i].x,2)+pow(vparr[i].y,2)),vparr[i].z,Bfphi));
                printf("gpu:%.10lf\n", valuearr[i]);
            }
        }
        // Main loop
        for(int i = 0; i < 40000; i++)
        {
            cudainterp<<<nblock,nthread>>>(cuvparr, cuBra, cuBrf);
            cudainterp<<<nblock,nthread>>>(cuvparr, cuBza, cuBzf);
            cudainterp<<<nblock,nthread>>>(cuvparr, cuBphia, cuBphif);
            cudavectorfieldform<<<nblock,nthread>>>(cuvBarr,cuvparr,cuBra,cuBza,cuBphia);
            cupush<<<nblock,nthread>>>(cuvparr,cuvvarr,cuvBarr, dt, qmr);
            cudamdp<<<nblock,nthread>>>(cuvparr,cuvvarr,cuvBarr,cuvgdc,custep, cufreq,i);
        }
        // transfer data from gpu to host
        cudaMemcpy(vparr, cuvparr, sizeof(vec)*nptc, cudaMemcpyDeviceToHost);
        cudaMemcpy(vvarr, cuvvarr, sizeof(vec)*nptc, cudaMemcpyDeviceToHost);
        cudaMemcpy(vfreq, cufreq, sizeof(double)*nptc, cudaMemcpyDeviceToHost);
        cudaMemcpy(custep, vstart, sizeof(int)*nptc, cudaMemcpyDeviceToHost);
        
        for (int i = 0; i < nptc; i += 4*655)
        {
            //diagnose(vparr[i], vvarr[i]);
            printf("freq:%lf\n", vfreq[i]);
            //printf("%s %d\n", mdpcross[i]?"true":"false", vtime[i]);  
        }
        int crosscount = 0;
        for(int i = 0; i < nptc; i++)
        {
            if (vfreq[i] != 0.0)
                crosscount++;
        }
        printf("\ncrosscount: %d\n", crosscount);
        output(vparr,vvarr,vfreq,1024,"parts.csv");

        cudaFree(cuBrf);
        cudaFree(cuBzf);
        cudaFree(cuBphif);
        cudaFree(cuvparr);
    }
    else
    {
    for (int k = 12*655; k < 32768; k += 104*655)
    {
        vec vv = vvarr[k];
        vec vp = vparr[k];
        
        int timestart = 0;
        double* gdcz = (double *)malloc(sizeof(double)*Ncycle);
        double zsumpre = 0.0;
        for(int i = 0; i < Ncycle; i++) gdcz[i] = 0.0;
        for(int i = 1; i < 100001; i++)
        {
            if(checkboundary(vp) == false)
            {
                printf("particle id %d outside boundary", k);
                break;
            }
            

            vec vB = vectorfield(vp);
            vec vs = dt*qmr/2.0*vB;
            double h2 = dot(vs, vs);
            vec vfor = cross(vv, vs);
            vec vback = cross(vfor,vs);
            vv = vv + (2.0/(1.0+h2))*(vfor+vback);
            vp = vp + dt*vv;     
            vec gdc = vp + 1.0/qmr/dot(vB,vB)*cross(vv, vB);
            gdcz[i%32] = gdc.z;
            if (i > 1000 && i % (Ncycle) == 0)
            {
                double zsum = 0.0;
                for(int p = 0; p < Ncycle; p++) zsum += gdcz[p];
                if (k == 12*655) printf("zsum: %lf\n", zsum);
                
                if (zsum*zsumpre < 0.0)
                {
                    //printf("crossing midplane time: %d\n",i);
                    if (timestart == 0) timestart = i;
                    else
                    {
                        printf("period:%.2f\n", (float)(i - timestart)/Ncycle);
                        //break;
                    }             
                }
                zsumpre = zsum;
            }
        }
        vec vB = vectorfield(vp);
        vec gdc = vp + 1.0/qmr/dot(vB,vB)*cross(vv, vB);
        double vpar2 = pow(dot(vv, vB),2)/dot(vB,vB);
        double vper2 = dot(cross(vv,vB),cross(vv,vB))/dot(vB,vB);
        double Bscalar = pow(dot(vB, vB),0.5);
        //printf("mu:%lf\n", vper2/Bscalar);
        //printf("vpar2:%lf vper2:%lf |B|:%lf pr:%lf\n",vpar2,vper2,Bscalar, pr);
        diagnose(vp, vv,'c');
        //printf("%lf, %lf, %lf, %lf, %lf, %lf\n", vp.x, vp.y, vp.z, vv.x, vv.y, vv.z);
        //printf("%lf, %lf, %lf\n", pr, pphi/(2*pi)*360.0, vp.z);
        if (&gdc != NULL)
        {
            //printf("%lf, %lf, %lf\n",gdc.x, gdc.y, gdc.z);
            double B1 = interp(sqrt(gdc.x*gdc.x+gdc.y*gdc.y),gdc.z,Bfr);
            double B2 = interp(sqrt(gdc.x*gdc.x+gdc.y*gdc.y),gdc.z,Bfz);
            double B3 = interp(sqrt(gdc.x*gdc.x+gdc.y*gdc.y),gdc.z,Bfphi);
            Bscalar = sqrt(B1*B1 + B2*B2 + B3*B3);
            //printf("guiding center mu:%lf ", vper2/Bscalar);
        }
        if (&gdc != NULL)
        {
            double gdcpsi = interp(sqrt(gdc.x*gdc.x+gdc.y*gdc.y),gdc.z,Apsi);
            double Pphi = -qmr*gdcpsi + sqrt(vpar2)*sqrt(gdc.x*gdc.x + gdc.y*gdc.y);
            //printf("Pphi = %lf\n", Pphi);
        }
    }
    }
    return 0;
}

